package com.study.crypto.certificate.server.advice;

import com.alibaba.fastjson.JSON;
import com.study.crypto.certificate.server.exception.SignatureException;
import com.study.crypto.dto.SignInfoDto;
import com.study.crypto.dto.ca.base.ResponseCaDto;
import com.study.crypto.general.spring.annotation.RequireSignature;
import com.study.crypto.signer.Signer;
import com.study.crypto.signer.SignerFactory;
import com.study.crypto.utils.KeyUtils;
import com.study.crypto.utils.ResourceUtils;
import lombok.extern.slf4j.Slf4j;
import org.bouncycastle.asn1.gm.GMObjectIdentifiers;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.crypto.CryptoException;
import org.bouncycastle.util.encoders.Base64;
import org.springframework.beans.BeanWrapperImpl;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.core.MethodParameter;
import org.springframework.core.annotation.AnnotatedElementUtils;
import org.springframework.core.annotation.Order;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import java.io.IOException;
import java.security.GeneralSecurityException;
import java.security.PrivateKey;

/**
 * 对响应数据进行签名，方法上使用了 @RequireSignature 注解
 * @author Songjin
 * @since 2022-11-10 13:39
 */
@Slf4j
@Order(1)
@ControllerAdvice
public class ResponseBodySignatureAdvice implements ResponseBodyAdvice<Object>, InitializingBean {

    /** 服务器私钥 */
    private PrivateKey privateKey;
    private final Signer signer = SignerFactory.produce(GMObjectIdentifiers.sm2sign_with_sm3);

    @Override
    public boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
        return (AnnotatedElementUtils.hasAnnotation(returnType.getContainingClass(), RequireSignature.class) ||
                returnType.hasMethodAnnotation(RequireSignature.class));
    }

    @Override
    public Object beforeBodyWrite(Object body, MethodParameter returnType, MediaType selectedContentType, Class<? extends HttpMessageConverter<?>> selectedConverterType, ServerHttpRequest request, ServerHttpResponse response) {
        try {
            ResponseCaDto responseBody = (ResponseCaDto) body;
            responseBody.setSignInfo(null);
            if (log.isInfoEnabled()) {
                log.info("对返回数据做签名，签待名数据: {}", JSON.toJSONString(responseBody));
            }
            // 对返回数据进行签名
            byte[] signature_ = signer.sign(JSON.toJSONBytes(body), privateKey);
            SignInfoDto signInfo = new SignInfoDto();
            signInfo.setSignAlgorithm(GMObjectIdentifiers.sm2sign_with_sm3.getId());
            signInfo.setSignValue(Base64.toBase64String(signature_));
            responseBody.setSignInfo(signInfo);
        } catch (IOException | CryptoException | GeneralSecurityException | CMSException e) {
            BeanWrapperImpl beanWrapper = new BeanWrapperImpl(body);
            beanWrapper.setPropertyValue("data", null);
            beanWrapper.setPropertyValue("resultCode", "-1");
            beanWrapper.setPropertyValue("resultCodeMsg", "签名异常");
            throw new SignatureException("签名异常", body, e);
        }
        return body;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        String pemString = ResourceUtils.readFileToString("certs/sm2-server-key.keystore");
        privateKey = KeyUtils.privateKey(pemString);
    }
}
